{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"validation_experiment.ipynb","provenance":[],"collapsed_sections":["lIyC4QdahqtO","RytAtJR0rrtT","_1GP7rVz2qs6"],"last_runtime":{"build_target":"//learning/deepmind/public/tools/ml_python:ml_notebook","kind":"private"},"toc_visible":true},"kernelspec":{"display_name":"Python 3","name":"python3"}},"cells":[{"cell_type":"markdown","metadata":{"id":"FxldYgrLH9Hx"},"source":["# Experiments reported in \"[Domain Conditional Predictors for Domain Adaptation](http://preregister.science/papers_20neurips/45_paper.pdf)\"\r\n","\r\n","Copyright 2021 Google LLC\r\n","\r\n","Licensed under the Apache License, Version 2.0 (the \"License\");\r\n","you may not use this file except in compliance with the License.\r\n","You may obtain a copy of the License at\r\n","\r\n","    https://www.apache.org/licenses/LICENSE-2.0\r\n","\r\n","Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."]},{"cell_type":"markdown","metadata":{"id":"4UlolfqgzFaP"},"source":["# Preamble"]},{"cell_type":"code","metadata":{"id":"_UxudP4jw_zc"},"source":["#@test {\"skip\": true}\r\n","!pip install dm-sonnet==2.0.0 --quiet\r\n","!pip install tensorflow_addons==0.12 --quiet"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"id":"NmYcx1XeIypw"},"source":["#@test {\"output\": \"ignore\"}\n","import numpy as np\n","import tensorflow as tf\n","import tensorflow_datasets as tfds\n","import tensorflow_addons as tfa\n","\n","try:\n","  import sonnet.v2 as snt\n","except ModuleNotFoundError:\n","  import sonnet as snt"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"o2OhDpEjvqVj"},"source":["#### Colab tested with:\r\n","\r\n","```\r\n","       TensorFlow version: 2.4.1\r\n","           Sonnet version: 2.0.0\r\n","TensorFlow Addons version: 0.12.0\r\n","```"]},{"cell_type":"code","metadata":{"id":"TvTi5uOmSGva"},"source":["#@test {\"skip\": true}\n","print(\"       TensorFlow version: {}\".format(tf.__version__))\n","print(\"           Sonnet version: {}\".format(snt.__version__))\n","print(\"TensorFlow Addons version: {}\".format(tfa.__version__))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"Ynxc0m8c9VRq"},"source":["# Data preparation\n","\n","Define 4 domains by transforming the data on the fly. Current transformations are rotation, blurring, flipping colors between background and digits, and horizontal flip."]},{"cell_type":"code","metadata":{"id":"_n4FjSeaSmfN"},"source":["#@test {\"output\": \"ignore\"}\n","batch_size = 100\n","NUM_DOMAINS = 4\n","\n","def process_batch_train(images, labels):\n","  images = tf.image.grayscale_to_rgb(images)\n","  images = tf.cast(images, dtype=tf.float32)\n","  images = images / 255.\n","  domain_index_candidates = tf.convert_to_tensor(list(range(NUM_DOMAINS)))\n","  samples = tf.random.categorical(tf.math.log([[1/NUM_DOMAINS for i in range(NUM_DOMAINS)]]), 1) # note log-prob\n","  domain_index=domain_index_candidates[tf.cast(samples[0][0], dtype=tf.int64)]\n","  if tf.math.equal(domain_index, tf.constant(0)):\n","    images = tfa.image.rotate(images, np.pi/3)\n","  elif tf.math.equal(domain_index, tf.constant(1)):\n","    images = tfa.image.gaussian_filter2d(images, filter_shape=[8,8])\n","  elif tf.math.equal(domain_index, tf.constant(2)):\n","    images = tf.ones_like(images) - images\n","  elif tf.math.equal(domain_index, tf.constant(3)):\n","    images = tf.image.flip_left_right(images)\n","  domain_label = tf.cast(domain_index, tf.int64)\n","  return images, labels, domain_label\n","\n","def process_batch_test(images, labels):\n","  images = tf.image.grayscale_to_rgb(images)\n","  images = tf.cast(images, dtype=tf.float32)\n","  images = images / 255.\n","  return images, labels\n","\n","def mnist(split, multi_domain_test=False):\n","  dataset = tfds.load(\"mnist\", split=split, as_supervised=True)\n","\n","  if split == \"train\":\n","    process_batch = process_batch_train\n","  else:\n","    if multi_domain_test:\n","      process_batch = process_batch_train\n","    else:\n","      process_batch = process_batch_test\n","\n","  dataset = dataset.map(process_batch)\n","  dataset = dataset.batch(batch_size)\n","  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)\n","  dataset = dataset.cache()\n","  return dataset\n","\n","mnist_train = mnist(\"train\").shuffle(1000)\n","mnist_test = mnist(\"test\")\n","mnist_test_multidomain = mnist(\"test\", multi_domain_test=True)"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"Ih4wtj1J90tY"},"source":["Look at samples from the training domains. Domain labels are such that: Rotation >> 0, Blurring >> 1, Color flipping >> 2, Horizontal flip >> 3."]},{"cell_type":"code","metadata":{"id":"iZX2KfjxTaJf"},"source":["#@test {\"skip\": true}\n","import matplotlib.pyplot as plt\n","images, label, domain_label = next(iter(mnist_train))\n","print(label[0], domain_label[0])\n","plt.imshow(images[0]);"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"3k0fh9-c-sEY"},"source":["# Baseline 1: Unconditional model\n","\n","A baseline model is defined below and referred to as unconditional since it does not take domain labels into account in any way."]},{"cell_type":"code","metadata":{"id":"JX44QcdyjCYF"},"source":["class M_unconditional(snt.Module):\n","\n","  def __init__(self):\n","    super(M_unconditional, self).__init__()\n","    self.hidden1 = snt.Conv2D(output_channels=10, kernel_shape=5, name=\"hidden1\")\n","    self.hidden2 = snt.Conv2D(output_channels=20, kernel_shape=5, name=\"hidden2\")\n","    self.flatten = snt.Flatten()\n","    self.logits = snt.Linear(10, name=\"logits\")\n","\n","  def __call__(self, images):\n","    output = tf.nn.relu(self.hidden1(images))\n","    output = tf.nn.relu(self.hidden2(output))\n","    output = self.flatten(output)\n","    output = self.logits(output)\n","    return output"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"id":"30B092fakLwC"},"source":["m_unconditional = M_unconditional()"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"9n_M8uKtqPZx"},"source":["Training of the baseline:"]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"b_7SUEKzkXVn"},"source":["#@test {\"output\": \"ignore\"}\n","opt_unconditional = snt.optimizers.SGD(learning_rate=0.01)\n","\n","num_epochs = 10\n","loss_log_unconditional = []\n","\n","def step(images, labels):\n","  \"\"\"Performs one optimizer step on a single mini-batch.\"\"\"\n","  with tf.GradientTape() as tape:\n","    logits_unconditional = m_unconditional(images)\n","    loss_unconditional = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits_unconditional,\n","                                                          labels=labels)\n","\n","    loss_unconditional = tf.reduce_mean(loss_unconditional)\n","\n","  params_unconditional = m_unconditional.trainable_variables\n","  grads_unconditional = tape.gradient(loss_unconditional, params_unconditional)\n","  opt_unconditional.apply(grads_unconditional, params_unconditional)\n","  return loss_unconditional\n","\n","for images, labels, domain_labels in mnist_train.repeat(num_epochs):\n","  loss_unconditional = step(images, labels)\n","  loss_log_unconditional.append(loss_unconditional.numpy())\n","\n","print(\"\\n\\nFinal loss: {}\".format(loss_log_unconditional[-1]))"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"id":"1oRsyRS2q4fw"},"source":["REDUCTION_FACTOR = 0.2  ## Factor in [0,1] used to check whether the training loss reduces during training\r\n","\r\n","## Checks whether the training loss reduces\r\n","assert loss_log_unconditional[-1] < REDUCTION_FACTOR*loss_log_unconditional[0]"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"RnClUVaGxg-y"},"source":["# Baseline 2: Domain invariant representations\n","\n","DANN-like model where the domain discriminator is replaced by a domain classifier aiming to induce invariance across training domains"]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"HlxrT9RHyJ--"},"source":["#@test {\"skip\": true}\n","class DANN_task(snt.Module):\n","\n","  def __init__(self):\n","    super(DANN_task, self).__init__()\n","    self.hidden1 = snt.Conv2D(output_channels=10, kernel_shape=5, name=\"hidden1\")\n","    self.hidden2 = snt.Conv2D(output_channels=20, kernel_shape=5, name=\"hidden2\")\n","    self.flatten = snt.Flatten()\n","    self.logits = snt.Linear(10, name=\"logits\")\n","\n","  def __call__(self, images):\n","    output = tf.nn.relu(self.hidden1(images))\n","    output = tf.nn.relu(self.hidden2(output))\n","    z = self.flatten(output)\n","    output = self.logits(z)\n","    return output, z"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"jRZQNjQ-yUNO"},"source":["#@test {\"skip\": true}\n","class DANN_domain(snt.Module):\n","\n","  def __init__(self):\n","    super(DANN_domain, self).__init__()\n","    self.logits = snt.Linear(NUM_DOMAINS, name=\"logits\")\n","\n","  def __call__(self, z):\n","    output = self.logits(z)\n","    return output"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"V5t0xjPiX9KX"},"source":["#@test {\"skip\": true}\n","m_DANN_task = DANN_task()\n","m_DANN_domain = DANN_domain()"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"QvC9oQPOysew"},"source":["Training of the DANN baseline"]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"3QfDa35Dyy5V"},"source":["#@test {\"skip\": true}\n","opt_task = snt.optimizers.SGD(learning_rate=0.01)\n","opt_domain = snt.optimizers.SGD(learning_rate=0.01)\n","domain_loss_weight = 0.2 ## Hyperparameter - factor to be multiplied by the domain entropy term when training the task classifier\n","\n","num_epochs = 20 ## Doubled the number of epochs to train the task classifier for as many iterations as the other methods since we have alternate updates\n","loss_log_dann = {'task_loss':[],'domain_loss':[]}\n","number_of_iterations = 0\n","\n","def step(images, labels, domain_labels, iteration_count):\n","  \"\"\"Performs one optimizer step on a single mini-batch.\"\"\"\n","  if iteration_count%2==0: ## Alternate between training the class classifier and the domain classifier\n","    with tf.GradientTape() as tape:\n","      logits_DANN_task, z_DANN = m_DANN_task(images)\n","      logist_DANN_domain = m_DANN_domain(z_DANN)\n","      loss_DANN_task = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits_DANN_task,\n","                                                            labels=labels)\n","      loss_DANN_domain = tf.nn.softmax_cross_entropy_with_logits(logits=logist_DANN_domain,\n","                                                            labels=1/NUM_DOMAINS*tf.ones_like(logist_DANN_domain)) ## Negative entropy of P(Y|X) measured as the cross-entropy against the uniform dist.\n","\n","      loss_DANN = tf.reduce_mean(loss_DANN_task + domain_loss_weight*loss_DANN_domain)\n","\n","    params_DANN = m_DANN_task.trainable_variables\n","    grads_DANN = tape.gradient(loss_DANN, params_DANN)\n","    opt_task.apply(grads_DANN, params_DANN)\n","    return 'task_loss', loss_DANN\n","  else:\n","    with tf.GradientTape() as tape:\n","      _, z_DANN = m_DANN_task(images)\n","      logist_DANN_domain_classifier = m_DANN_domain(z_DANN)\n","      loss_DANN_domain_classifier = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logist_DANN_domain_classifier,\n","                                                            labels=domain_labels)\n","\n","      loss_DANN_domain_classifier = tf.reduce_mean(loss_DANN_domain_classifier)\n","\n","    params_DANN_domain_classifier = m_DANN_domain.trainable_variables\n","    grads_DANN_domain_classifier = tape.gradient(loss_DANN_domain_classifier, params_DANN_domain_classifier)\n","    opt_domain.apply(grads_DANN_domain_classifier, params_DANN_domain_classifier)\n","    return 'domain_loss', loss_DANN_domain_classifier\n","\n","for images, labels, domain_labels in mnist_train.repeat(num_epochs):\n","  number_of_iterations += 1\n","  loss_tag, loss_dann = step(images, labels, domain_labels, number_of_iterations)\n","  loss_log_dann[loss_tag].append(loss_dann.numpy())\n","\n","print(\"\\n\\nFinal losses: {} - {}, {} - {}\".format('task_loss', loss_log_dann['task_loss'][-1], 'domain_loss', loss_log_dann['domain_loss'][-1]))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"AIkz-7FTqVl-"},"source":["# Definition of our models\n","\n","The models for our proposed setting are defined below. \n","\n","*   The FiLM layer simply projects z onto 2 tensors (independent dense layers for each projection) matching the shape of the features. Each such tensor is used for element-wise multiplication and addition with the input features.\n","*   m_domain corresponds to a domain classifier. It outputs the output of the second conv. layer to be used as z, as well as a set of logits over the set of train domains.\n","*   m_task is the main classifier and it contains FiLM layers that take z as input. Its output corresponds to the set of logits over the labels."]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"Xq2Wh-fBxi6L"},"source":["#@test {\"skip\": true}\n","class FiLM(snt.Module):\n","  def __init__(self, features_shape):\n","    super(FiLM, self).__init__()\n","    self.features_shape = features_shape\n","    target_dimension = np.prod(features_shape)\n","    self.hidden_W = snt.Linear(output_size=target_dimension)\n","    self.hidden_B = snt.Linear(output_size=target_dimension)\n","\n","  def __call__(self, features, z):\n","    W = snt.reshape(self.hidden_W(z), output_shape=self.features_shape)\n","    B = snt.reshape(self.hidden_B(z), output_shape=self.features_shape)\n","    output = W*features+B\n","    return output"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"bFWGf78Gy4yf"},"source":["#@test {\"skip\": true}\n","class M_task(snt.Module):\n","\n","  def __init__(self):\n","    super(M_task, self).__init__()\n","    self.hidden1 = snt.Conv2D(output_channels=10, kernel_shape=5, name=\"hidden1\")\n","    self.film1 = FiLM(features_shape=[28,28,10])\n","    self.hidden2 = snt.Conv2D(output_channels=20, kernel_shape=5, name=\"hidden2\")\n","    self.film2 = FiLM(features_shape=[28,28,20])\n","    self.flatten = snt.Flatten()\n","    self.logits = snt.Linear(10, name=\"logits\")\n","\n","  def __call__(self, images, z):\n","    output = tf.nn.relu(self.hidden1(images))\n","    output = self.film1(output,z)\n","    output = tf.nn.relu(self.hidden2(output))\n","    output = self.film2(output,z)\n","    output = self.flatten(output)\n","    output = self.logits(output)\n","    return output"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"TPZWVODA41Vu"},"source":["#@test {\"skip\": true}\n","class M_domain(snt.Module):\n","\n","  def __init__(self):\n","    super(M_domain, self).__init__()\n","    self.hidden = snt.Conv2D(output_channels=10, kernel_shape=5, name=\"hidden\")\n","    self.flatten = snt.Flatten()\n","    self.logits = snt.Linear(NUM_DOMAINS, name=\"logits\")\n","\n","  def __call__(self, images):\n","    output = tf.nn.relu(self.hidden(images))\n","    z = self.flatten(output)\n","    output = self.logits(z)\n","    return output, z"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"ASPB5gaQT2eP"},"source":["#@test {\"skip\": true}\n","m_task = M_task()\n","m_domain = M_domain()"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"lbHGLMsmT8Cu"},"source":["#@test {\"skip\": true}\n","images, labels = next(iter(mnist_test))\n","domain_logits, z = m_domain(images)\n","logits = m_task(images, z)\n","  \n","prediction = tf.argmax(logits[0]).numpy()\n","actual = labels[0].numpy()\n","print(\"Predicted class: {} actual class: {}\".format(prediction, actual))\n","plt.imshow(images[0])"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"UA7qo9jnrdey"},"source":["# Training of the proposed model"]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"PeT7ssTLiEAG"},"source":["#@test {\"skip\": true}\n","from tqdm import tqdm\n","\n","# MNIST training set has 60k images.\n","num_images = 60000\n","\n","def progress_bar(generator):\n","  return tqdm(\n","      generator,\n","      unit='images',\n","      unit_scale=batch_size,\n","      total=(num_images // batch_size) * num_epochs)"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"f10HudmRUDzN"},"source":["#@test {\"skip\": true}\n","opt = snt.optimizers.SGD(learning_rate=0.01)\n","\n","num_epochs = 10\n","\n","loss_log = {'total_loss':[], 'task_loss':[], 'domain_loss':[]}\n","\n","def step(images, labels, domain_labels):\n","  \"\"\"Performs one optimizer step on a single mini-batch.\"\"\"\n","  with tf.GradientTape() as tape:\n","    domain_logits, z = m_domain(images)\n","    logits = m_task(images, z)\n","    loss_task = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,\n","                                                          labels=labels)\n","    loss_domain = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=domain_logits,\n","                                                          labels=domain_labels)\n","\n","    loss = loss_task + loss_domain\n","\n","    loss = tf.reduce_mean(loss)\n","    loss_task = tf.reduce_mean(loss_task)\n","    loss_domain = tf.reduce_mean(loss_domain)\n","\n","  params = m_task.trainable_variables + m_domain.trainable_variables\n","  grads = tape.gradient(loss, params)\n","  opt.apply(grads, params)\n","  return loss, loss_task, loss_domain\n","\n","for images, labels, domain_labels in progress_bar(mnist_train.repeat(num_epochs)):\n","  loss, loss_task, loss_domain = step(images, labels, domain_labels)\n","  loss_log['total_loss'].append(loss.numpy())\n","  loss_log['task_loss'].append(loss_task.numpy())\n","  loss_log['domain_loss'].append(loss_domain.numpy())\n","\n","print(\"\\n\\nFinal total loss: {}\".format(loss.numpy()))\n","print(\"\\n\\nFinal task loss: {}\".format(loss_task.numpy()))\n","print(\"\\n\\nFinal domain loss: {}\".format(loss_domain.numpy()))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"lbgKFbL096ti"},"source":["# Ablation 1: Learned domain-wise context variable z\n","\n","Here we consider a case where the context variables z used for conditioning are learned directly from data, and the domain predictor is discarded. **This only allows for in-domain prediction though**."]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"URp7LAKe-nc-"},"source":["#@test {\"skip\": true}\n","class M_learned_z(snt.Module):\n","\n","  def __init__(self):\n","    super(M_learned_z, self).__init__()\n","    self.context = snt.Embed(vocab_size=NUM_DOMAINS, embed_dim=128)\n","    self.hidden1 = snt.Conv2D(output_channels=10, kernel_shape=5, name=\"hidden1\")\n","    self.film1 = FiLM(features_shape=[28,28,10])\n","    self.hidden2 = snt.Conv2D(output_channels=20, kernel_shape=5, name=\"hidden2\")\n","    self.film2 = FiLM(features_shape=[28,28,20])\n","    self.flatten = snt.Flatten()\n","    self.logits = snt.Linear(10, name=\"logits\")\n","\n","  def __call__(self, images, domain_labels):\n","    z = self.context(domain_labels)\n","    output = tf.nn.relu(self.hidden1(images))\n","    output = self.film1(output,z)\n","    output = tf.nn.relu(self.hidden2(output))\n","    output = self.film2(output,z)\n","    output = self.flatten(output)\n","    output = self.logits(output)\n","    return output"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"_gP0ge2MDYjn"},"source":["#@test {\"skip\": true}\r\n","m_learned_z = M_learned_z()"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"vbQ6SpviFIjM"},"source":["#@test {\"skip\": true}\n","opt_learned_z = snt.optimizers.SGD(learning_rate=0.01)\n","\n","num_epochs = 10\n","loss_log_learned_z = []\n","\n","def step(images, labels, domain_labels):\n","  \"\"\"Performs one optimizer step on a single mini-batch.\"\"\"\n","  with tf.GradientTape() as tape:\n","    logits_learned_z = m_learned_z(images, domain_labels)\n","    loss_learned_z = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits_learned_z,\n","                                                          labels=labels)\n","\n","    loss_learned_z = tf.reduce_mean(loss_learned_z)\n","\n","  params_learned_z = m_learned_z.trainable_variables\n","  grads_learned_z = tape.gradient(loss_learned_z, params_learned_z)\n","  opt_learned_z.apply(grads_learned_z, params_learned_z)\n","  return loss_learned_z\n","\n","for images, labels, domain_labels in mnist_train.repeat(num_epochs):\n","  loss_learned_z = step(images, labels, domain_labels)\n","  loss_log_learned_z.append(loss_learned_z.numpy())\n","\n","print(\"\\n\\nFinal loss: {}\".format(loss_log_learned_z[-1]))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"Y1QCC0soeEs2"},"source":["# Ablation 2: Dropping the domain classification term of the loss\n","\n","We consider an ablation where the same models as in our conditional predictor are used, but training is carried out with the classification loss only. This gives us a model with the same capacity as ours but no explicit mechanism to account for domain variations in train data. The goal of this ablation is to understand whether the improvement might be simply coming from the added capacity rather than the conditional modeling."]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"bX3dA9qKerpg"},"source":["#@test {\"skip\": true}\n","m_task_ablation = M_task()\n","m_domain_ablation = M_domain()\n","m_DANN_ablation = DANN_domain() ## Used for evaluating how domain dependent the representations are"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"S1zYXgQaey5b"},"source":["#@test {\"skip\": true}\n","opt_ablation = snt.optimizers.SGD(learning_rate=0.01)\n","\n","num_epochs = 10\n","\n","loss_log_ablation = []\n","\n","def step(images, labels, domain_labels):\n","  \"\"\"Performs one optimizer step on a single mini-batch.\"\"\"\n","  with tf.GradientTape() as tape:\n","    domain_logits_ablation, z = m_domain_ablation(images)\n","    logits_ablation = m_task_ablation(images, z)\n","    loss_ablation = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits_ablation,\n","                                                          labels=labels)\n","\n","    loss_ablation = tf.reduce_mean(loss_ablation)\n","\n","  params_ablation = m_task_ablation.trainable_variables + m_domain_ablation.trainable_variables\n","  grads_ablation = tape.gradient(loss_ablation, params_ablation)\n","  opt_ablation.apply(grads_ablation, params_ablation)\n","  return loss_ablation\n","\n","for images, labels, domain_labels in mnist_train.repeat(num_epochs):\n","  loss_ablation = step(images, labels, domain_labels)\n","  loss_log_ablation.append(loss_ablation.numpy())\n","\n","print(\"\\n\\nFinal task loss: {}\".format(loss_ablation.numpy()))"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"cIYqjnPU73dq"},"source":["#@test {\"skip\": true}\n","opt_ablation_domain_classifier = snt.optimizers.SGD(learning_rate=0.01)\n","\n","num_epochs = 10\n","\n","log_loss_ablation_domain_classification = []\n","\n","def step(images, labels, domain_labels):\n","  \"\"\"Performs one optimizer step on a single mini-batch.\"\"\"\n","  with tf.GradientTape() as tape:\n","    _, z = m_domain_ablation(images)\n","    logits_ablation_domain_classification = m_DANN_ablation(z)\n","    loss_ablation_domain_classification = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits_ablation_domain_classification,\n","                                                          labels=domain_labels)\n","\n","    loss_ablation_domain_classification = tf.reduce_mean(loss_ablation_domain_classification)\n","\n","  params_ablation_domain_classification = m_DANN_ablation.trainable_variables\n","  grads_ablation_domain_classification = tape.gradient(loss_ablation_domain_classification, params_ablation_domain_classification)\n","  opt_ablation.apply(grads_ablation_domain_classification, params_ablation_domain_classification)\n","  return loss_ablation_domain_classification\n","\n","for images, labels, domain_labels in mnist_train.repeat(num_epochs):\n","  loss_ablation_domain_classifier = step(images, labels, domain_labels)\n","  log_loss_ablation_domain_classification.append(loss_ablation_domain_classifier.numpy())\n","\n","print(\"\\n\\nFinal task loss: {}\".format(loss_ablation_domain_classifier.numpy()))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"EKTWrQcG2ibk"},"source":["# Results"]},{"cell_type":"markdown","metadata":{"id":"lIyC4QdahqtO"},"source":["## Plots of training losses"]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"7oHgkFoyx9sX"},"source":["#@test {\"skip\": true}\n","f = plt.figure(figsize=(32,8))\n","ax = f.add_subplot(1,3,1)\n","ax.plot(loss_log['total_loss'])\n","ax.set_title('Total Loss')\n","\n","ax = f.add_subplot(1,3,2)\n","ax.plot(loss_log['task_loss'])\n","ax.set_title('Task loss')\n","\n","ax = f.add_subplot(1,3,3)\n","ax.plot(loss_log['domain_loss'])\n","ax.set_title('Domain loss')\n","\n"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"o5q3CEB_GTLk"},"source":["#@test {\"skip\": true}\n","f = plt.figure(figsize=(8,8))\n","ax = f.add_axes([1,1,1,1])\n","ax.plot(loss_log_unconditional)\n","ax.set_title('Unconditional baseline - Train Loss')"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"Z8M2z8FiZGmm"},"source":["#@test {\"skip\": true}\n","f = plt.figure(figsize=(16,8))\n","ax = f.add_subplot(1,2,1)\n","ax.plot(loss_log_dann['task_loss'])\n","ax.set_title('Domain invariant baseline - Task loss (Class. + -Entropy)')\n","\n","ax = f.add_subplot(1,2,2)\n","ax.plot(loss_log_dann['domain_loss'])\n","ax.set_title('Domain invariant baseline - Domain classification loss')"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"7-dDIhoDbtNW"},"source":["#@test {\"skip\": true}\n","f = plt.figure(figsize=(16,8))\n","ax = f.add_subplot(1,2,1)\n","ax.plot(loss_log_learned_z)\n","ax.set_title('Ablation 1 - Task loss')\n","\n","ax = f.add_subplot(1,2,2)\n","ax.plot(loss_log_ablation)\n","ax.set_title('Ablation 2 - Task loss')"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"Mpms741T8wqC"},"source":["#@test {\"skip\": true}\n","f = plt.figure(figsize=(8,8))\n","ax = f.add_axes([1,1,1,1])\n","ax.plot(log_loss_ablation_domain_classification)\n","ax.set_title('Ablation 2: Domain classification - Train Loss')"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"7bNKWV7Grh72"},"source":["## Out-of-domain evaluations\n","\n","The original test set of mnist without any transformations is considered"]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"j7I0c0gki2Wp"},"source":["#@test {\"skip\": true}\n","total = 0\n","correct = 0\n","correct_unconditional = 0\n","correct_adversarial = 0\n","correct_ablation2 = 0 ## The model corresponding to ablation 1 can only be used with in-domain data (with domain labels)\n","for images, labels in mnist_test:\n","  domain_logits, z = m_domain(images)\n","  logits = m_task(images, z)\n","  logits_unconditional = m_unconditional(images)\n","  logits_adversarial, _ = m_DANN_task(images)\n","  domain_logits_ablation, z_ablation = m_domain_ablation(images)\n","  logits_ablation2 = m_task_ablation(images, z_ablation)\n","  predictions = tf.argmax(logits, axis=1)\n","  predictions_unconditional = tf.argmax(logits_unconditional, axis=1)\n","  predictions_adversarial = tf.argmax(logits_adversarial, axis=1)\n","  predictions_ablation2 = tf.argmax(logits_ablation2, axis=1)\n","  correct += tf.math.count_nonzero(tf.equal(predictions, labels))\n","  correct_unconditional += tf.math.count_nonzero(tf.equal(predictions_unconditional, labels))\n","  correct_adversarial += tf.math.count_nonzero(tf.equal(predictions_adversarial, labels))\n","  correct_ablation2 += tf.math.count_nonzero(tf.equal(predictions_ablation2, labels))\n","  total += images.shape[0]\n","\n","print(\"Got %d/%d (%.02f%%) correct\" % (correct, total, correct / total * 100.))\n","print(\"Unconditional baseline perf.: %d/%d (%.02f%%) correct\" % (correct_unconditional, total, correct_unconditional / total * 100.))\n","print(\"Adversarial baseline perf.: %d/%d (%.02f%%) correct\" % (correct_adversarial, total, correct_adversarial / total * 100.))\n","print(\"Ablation 2 perf.: %d/%d (%.02f%%) correct\" % (correct_ablation2, total, correct_ablation2 / total * 100.))"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"RytAtJR0rrtT"},"source":["## In-domain evaluations and domain prediction\n","\n","The same transformations applied in train data are applied during test"]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"uyTrvvQnOj-J"},"source":["#@test {\"skip\": true}\n","n_repetitions = 10 ## Going over the test set multiple times to account for multiple transformations\n","total = 0\n","correct_class = 0\n","correct_unconditional = 0\n","correct_adversarial = 0\n","correct_ablation1 = 0\n","correct_ablation2 = 0\n","correct_domain = 0\n","correct_domain_adversarial = 0\n","correct_domain_ablation = 0\n","for images, labels, domain_labels in mnist_test_multidomain.repeat(n_repetitions):\n","  domain_logits, z = m_domain(images)\n","  class_logits = m_task(images, z)\n","  logits_unconditional = m_unconditional(images)\n","  logits_adversarial, z_adversarial = m_DANN_task(images)\n","  domain_logits_adversarial = m_DANN_domain(z_adversarial)\n","  logits_ablation1 = m_learned_z(images, domain_labels)\n","  _, z_ablation = m_domain_ablation(images)\n","  domain_logits_ablation = m_DANN_ablation(z_ablation)\n","  logits_ablation2 = m_task_ablation(images, z_ablation)\n","\n","  predictions_class = tf.argmax(class_logits, axis=1)\n","  predictions_unconditional = tf.argmax(logits_unconditional, axis=1)\n","  predictions_adversarial = tf.argmax(logits_adversarial, axis=1)\n","  predictions_ablation1 = tf.argmax(logits_ablation1, axis=1)\n","  predictions_ablation2 = tf.argmax(logits_ablation2, axis=1)\n","  predictions_domain = tf.argmax(domain_logits, axis=1)\n","  predictions_domain_adversarial = tf.argmax(domain_logits_adversarial, axis=1)\n","  predictions_domain_ablation = tf.argmax(domain_logits_ablation, axis=1)\n","\n","  correct_class += tf.math.count_nonzero(tf.equal(predictions_class, labels))\n","  correct_unconditional += tf.math.count_nonzero(tf.equal(predictions_unconditional, labels))\n","  correct_adversarial += tf.math.count_nonzero(tf.equal(predictions_adversarial, labels))\n","  correct_ablation1 += tf.math.count_nonzero(tf.equal(predictions_ablation1, labels))\n","  correct_ablation2 += tf.math.count_nonzero(tf.equal(predictions_ablation2, labels))\n","  correct_domain += tf.math.count_nonzero(tf.equal(predictions_domain, domain_labels))\n","  correct_domain_adversarial += tf.math.count_nonzero(tf.equal(predictions_domain_adversarial, domain_labels))\n","  correct_domain_ablation += tf.math.count_nonzero(tf.equal(predictions_domain_ablation, domain_labels))\n","  total += images.shape[0]\n","\n","print(\"In domain unconditional baseline perf.: %d/%d (%.02f%%) correct\" % (correct_unconditional, total, correct_unconditional / total * 100.))\n","print(\"In domain adversarial baseline perf.: %d/%d (%.02f%%) correct\" % (correct_adversarial, total, correct_adversarial / total * 100.))\n","print(\"In domain ablation 1: %d/%d (%.02f%%) correct\" % (correct_ablation1, total, correct_ablation1 / total * 100.))\n","print(\"In domain ablation 2: %d/%d (%.02f%%) correct\" % (correct_ablation2, total, correct_ablation2 / total * 100.))\n","print(\"In domain class predictions: Got %d/%d (%.02f%%) correct\" % (correct_class, total, correct_class / total * 100.))\n","print(\"\\n\\nDomain predictions: Got %d/%d (%.02f%%) correct\" % (correct_domain, total, correct_domain / total * 100.))\n","print(\"Adversarial baseline domain predictions: Got %d/%d (%.02f%%) correct\" % (correct_domain_adversarial, total, correct_domain_adversarial / total * 100.))\n","print(\"Ablation 2 domain predictions: Got %d/%d (%.02f%%) correct\" % (correct_domain_ablation, total, correct_domain_ablation / total * 100.))"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"CpxOFsBviqsb"},"source":["#@test {\"skip\": true}\n","def sample(correct, rows, cols):\n","  n = 0\n","\n","  f, ax = plt.subplots(rows, cols)\n","  if rows > 1:    \n","    ax = tf.nest.flatten([tuple(ax[i]) for i in range(rows)])\n","  f.set_figwidth(14)\n","  f.set_figheight(4 * rows)\n","\n","\n","  for images, labels in mnist_test:\n","    domain_logits, z = m_domain(images)\n","    logits = m_task(images, z)\n","    predictions = tf.argmax(logits, axis=1)\n","    eq = tf.equal(predictions, labels)\n","    for i, x in enumerate(eq):\n","      if x.numpy() == correct:\n","        label = labels[i]\n","        prediction = predictions[i]\n","        image = images[i]\n","\n","        ax[n].imshow(image)\n","        ax[n].set_title(\"Prediction:{}\\nActual:{}\".format(prediction, label))\n","\n","        n += 1\n","        if n == (rows * cols):\n","          break\n","\n","    if n == (rows * cols):\n","      break"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"_1GP7rVz2qs6"},"source":["## Samples and corresponding predictions"]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"4wrojuocUPjy"},"source":["#@test {\"skip\": true}\r\n","sample(correct=True, rows=1, cols=5)"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"background_save":true},"id":"lbBKFSiRaLwc"},"source":["#@test {\"skip\": true}\r\n","sample(correct=False, rows=2, cols=5)"],"execution_count":null,"outputs":[]}]}